import os
import common.ReadFile as rf
import tensorflow as tf


# import matplotlib.pyplot as plt

def _parse_function(example_protocol):
    features = {"label": tf.FixedLenFeature((), tf.int64, default_value=0),
                "image": tf.FixedLenFeature((), tf.string, default_value="")}
    # features = {"label": tf.FixedLenFeature((), tf.int64, default_value=0)}
    parsed_features = tf.parse_single_example(example_protocol, features)
    image = tf.decode_raw(parsed_features["image"], tf.uint8)
    image = tf.reshape(image, [64, 64, 1])
    image = tf.cast(image, tf.float32) / 255.0
    label = tf.cast(parsed_features["label"], tf.int32)
    return image, label
    # return parsed_features["label"]


current_directory = os.getcwd()
end = (len("model") + 1) * (-1)
train_directory = current_directory[0:end] + '/data/tf_train_data_part'
test_directory = current_directory[0:end] + '/data/tf_test_data_part'
tail = 'tfrecords'

# train data
read_file = rf.ReadFile(train_directory, tail)
read_file.run()
dataset = tf.data.TFRecordDataset(read_file.full_path)
dataset = dataset.map(_parse_function)
dataset = dataset.batch(140).repeat()
# iterator = dataset.make_one_shot_iterator()
# image, label = iterator.get_next()
# dataset = dataset.repeat()

# # test data
# read_file_test = rf.ReadFile(test_directory, tail)
# read_file_test.run()
# validate_dataset = tf.data.TFRecordDataset(read_file_test.full_path)
# validate_dataset = validate_dataset.map(_parse_function)
# validate_dataset = validate_dataset.batch(3755)
# validate_dataset = validate_dataset.repeat()

# user tf.keras
# print (tf.keras.__version__)
model = tf.keras.Sequential()
model.add(tf.keras.layers.Conv2D(32, (5, 5), (1, 1), 'same', activation='relu'))
model.add(tf.keras.layers.MaxPooling2D((2, 2), (1, 1), 'same'))
model.add(tf.keras.layers.Conv2D(64, (5, 5), (1, 1), 'same', activation='relu'))
model.add(tf.keras.layers.MaxPooling2D((2, 2), (1, 1), 'same'))
model.add(tf.keras.layers.Flatten())
model.add(tf.keras.layers.Dense(1024, activation=tf.nn.relu))
model.add(tf.keras.layers.Dropout(0.5))
model.add(tf.keras.layers.Dense(140, activation=tf.nn.softmax))

model.compile(optimizer=tf.train.AdamOptimizer(1e-4),
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(dataset, epochs=50, steps_per_epoch=160)

model.save_weights('mymodel.h5', save_format='h5')

# show the pic
# with tf.Session() as sess:
#     plt.figure()
#     for i in range(2):
#         labels, images = sess.run([label, image])
#         plt.subplot(1, 2, i+1)
#         plt.title(labels)
#         plt.imshow(images, cmap='gray')
#     plt.show()

# user tf.sess
